#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
模型工具类
提供模型路径查找等公共功能
"""

import os


def find_local_model_path(project_root=None):
    """
    查找本地模型路径
    
    Args:
        project_root (str): 项目根目录路径，默认为None时使用当前文件所在目录的上级目录
        
    Returns:
        str or None: 找到的模型路径，如果没找到则返回None
    """
    if project_root is None:
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    
    # 项目内模型路径
    project_model_path = os.path.join(project_root, "models", "Qwen1.5-1.8B")
    
    # Hugging Face缓存模型路径
    hf_model_path = os.path.expanduser("~/.cache/huggingface/hub/models--Qwen--Qwen1.5-1.8B")
    
    local_model_path = None
    
    # 查找项目内的模型
    if os.path.exists(project_model_path):
        # 检查是否是直接的模型文件（不是Hugging Face的结构）
        config_path = os.path.join(project_model_path, "config.json")
        if os.path.exists(config_path):
            local_model_path = project_model_path
    
    # 如果项目内没有模型，则查找Hugging Face缓存中的模型
    if local_model_path is None and os.path.exists(hf_model_path):
        refs_path = os.path.join(hf_model_path, "refs", "main")
        if os.path.exists(refs_path):
            with open(refs_path, 'r') as f:
                commit_hash = f.read().strip()
                snapshot_path = os.path.join(hf_model_path, "snapshots", commit_hash)
                if os.path.exists(snapshot_path):
                    local_model_path = snapshot_path
                    
    return local_model_path